#! /usr/bin/env python
#
# Pythonically read and write a FITS file using numpy arrays

import sys
import os

import numpy

import ctypes
from fitsio import cfitsio, pyfitsio

def get_test_data_3():
    '''3D test data'''

    nz = 10
    ny = 300
    nx = 200
    data = numpy.zeros( (nz, ny, nx), dtype=numpy.int16)
    
    for z in range(nz):
        for y in range(ny):
            for x in range(nx):
                data[z][y][x] = x + y + z

    return data

def get_test_data_4():
    '''4D test data'''

    nt = 20
    nz = 10
    ny = 256
    nx = 256
    data = numpy.zeros( (nt, nz, ny, nx), dtype=numpy.float32)
    
    for t in range(nt):
        for z in range(nz):
            for y in range(ny):
                for x in range(nx):
                    data[t][z][y][x] = t + x + y + z

    return data

def get_test_data_5():
    '''5D test data'''

    ns = 10
    nt = 20
    nz = 10
    ny = 32
    nx = 32
    data = numpy.zeros( (ns, nt, nz, ny, nx), dtype=numpy.float32)
    
    for s in range(ns):
        for t in range(nt):
            for z in range(nz):
                for y in range(ny):
                    for x in range(nx):
                        data[s][t][z][y][x] = s + t + x + y + z

    return data

def check_type(dtype):

    if dtype == numpy.int16:
        ttt = cfitsio.TSHORT, cfitsio.SHORT_IMG, ctypes.c_int16
    elif dtype == numpy.uint16:
        ttt = cfitsio.USHORT, cfitsio.USHORT_IMG, ctypes.c_uint16
    elif dtype == numpy.int32:
        ttt = cfitsio.TINT32BIT, cfitsio.LONG_IMG, ctypes.c_int32
    elif dtype == numpy.int64:
        ttt = cfitsio.TLONGLONG, cfitsio.LONGLONG_IMG, ctypes.c_int64
    elif dtype == numpy.float32:
        ttt = cfitsio.TFLOAT, cfitsio.FLOAT_IMG, ctypes.c_float
    elif dtype == numpy.float64:
        ttt = cfitsio.TDOUBLE, cfitsio.DOUBLE_IMG, ctypes.c_double
    else:
        print 'Unsupported numpy type by this example', dtype
        sys.exit(1)

    return ttt

def fits_to_numpy(ff, pixdata):

    ttype, itype, ctype = check_type(pixdata.dtype)
    
    pixptr = pixdata.ctypes.data_as(ctypes.POINTER(ctype))

    npix = pixdata.size
    ndim = len(pixdata.shape)

    fpixel = (ctypes.c_longlong * ndim)()
    for n in range(ndim):
        fpixel[n] = 1

    nelements = ctypes.c_longlong(npix)
    nullval = ctypes.c_float(0)
    anynull = ctypes.c_int()
    nbuffer = ctypes.c_longlong(npix)
    status  = ctypes.c_int()

    if cfitsio.fits_read_pix(ff.fptr, ttype, fpixel, nbuffer,
                             ctypes.byref(nullval), pixptr,
                             ctypes.byref(anynull), ctypes.byref(status)):

        pyfitsio.handle_error(status)

def numpy_to_fits(pixdata, ff):
    
    ttype, itype, ctype = check_type(pixdata.dtype)
    npix = pixdata.size
    ndim = len(pixdata.shape)

    fpixel = (ctypes.c_longlong * ndim)()
    for n in range(ndim):
        fpixel[n] = 1

    nelements = ctypes.c_longlong(npix)
    pixptr = pixdata.ctypes.data_as(ctypes.POINTER(ctype))
    status  = ctypes.c_int()

    if cfitsio.fits_write_pix(ff.fptr, ttype, fpixel, nelements, pixptr,
                              ctypes.byref(status)):
        pyfitsio.handle_error(status)

def wdata(filename):

    # Generate some data and put in a multidimensional numpy array. Just edit 
    # this line to call any of the other test data generators
    pixdata = get_test_data_3()
    print 'Got data', pixdata.shape, pixdata.size

    # Axes in the CFITSIO order. 
    naxes = list(pixdata.shape)

    # numpy uses row-major ordering (..,z, y, x), but CFITSIO uses column-major ordering (x, y, z, ...) so reverse the axes
    naxes.reverse()
    ndim = len(naxes)
    
    # Create new file for writing. Overwrite if already exists
    ff = pyfitsio.open(filename, mode='w', clobber=True)

    # Write header for a 3-d image with the dimensions of the numpy array
    ff.create_image( naxes, cfitsio.FLOAT_IMG)

    # Write some more keywords and comments
    ff.write_comment(' ')
    ff.write_comment('Demo file to show how to write ' + str(ndim) + ' dimensional data')
    ff.write_comment('File written by: fitsnp.py')
    ff.write_key('NAME', 'Ian Bond', 'Author of this demo program')

    # Write the data
    numpy_to_fits(pixdata, ff)

    ff.close()

    print 'Written to', filename

def rdata(filename):
    
    ff = pyfitsio.open(filename, 'r')

    # Get the axis dimensions in reverse order of CFITSIO 
    naxes = list(ff.read_axes())
    naxes.reverse()
    print filename, naxes

    pixdata = numpy.zeros(naxes, dtype=numpy.int16)

    fits_to_numpy(ff, pixdata)

    ff.close()

def main(filename):

    try:
        wdata(filename)
        #rdata(filename)
    except pyfitsio.FitsIOError, e:
        print 'Problem calling fitsio'
        print str(e)


if __name__ == '__main__':

    filename = 'numpytest.fit'
    main(filename)
